import asyncio
import logging
from contextlib import AsyncExitStack
from typing import Optional, Any

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


class PluginClient:
    def __init__(self, plugin_path: str):
        self.exit_stack = None
        self.write = None
        self.stdio = None
        self.session: Optional[ClientSession] = None
        self.plugin_path = plugin_path

    async def connect(self):
        self.exit_stack = AsyncExitStack()
        return await self.connect_to_server(self.plugin_path)

    async def connect_to_server(self, server_script_path: str):
        """Connect to an MCP server
        Args:
            server_script_path: Path to the server script (.py or .js)
        """
        logging.info(f'Start load plugin_path: {self.plugin_path}')
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')
        if not (is_python or is_js):
            raise ValueError("Server script must be a .py or .js file")

        command = "python" if is_python else "node"
        server_params = StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None
        )

        self.stdio, self.write = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))

        async_initialize = lambda: self.session.initialize()
        try:
            await asyncio.wait_for(async_initialize(), 10)
        except asyncio.TimeoutError:
            logging.error(
                f'Timeout to connect to MCP server, path: {server_script_path}')
        logging.info(f'Successfully load plugin_path: {self.plugin_path}')

    async def exec(self, tool_name: str, tool_args) -> Any:
        return await self.session.call_tool(tool_name, tool_args)

    async def cleanup(self):
        """Clean up resources"""
        await self.exit_stack.aclose()
